# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/5 19:31 
@Author  : ZhangShenao 
@File    : 使用智谱AI的GLM-4-0520模型实现文档翻译.py 
@Desc    : 使用智谱AI的GLM-4-0520模型实现文档翻译
"""
import os
from typing import List

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.language_models import BaseChatModel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# 加载环境变量
dotenv.load_dotenv()

# 构造Prompt
prompt = ChatPromptTemplate.from_messages([
    ("system",
     "你是一个智能翻译专家，请将用户输入的唐诗，生成{language}的准确翻译，要求做到信达雅，尽可能保留原文的韵味和意境。"),
    ("human", "{poem}")
])

# 创建多个LLM进行对比
llms: List[BaseChatModel] = [
    ChatZhipuAI(model="glm-4-0520", temperature=0),
    ChatZhipuAI(model="glm-4-airx", temperature=0),
    ChatOpenAI(model="gpt-4o-mini", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
]
parser = StrOutputParser()

poem = """床前明月光，
疑似地上霜。
举头望明月，
低头思故乡。"""

for llm in llms:
    chain = prompt | llm | parser
    translation = chain.invoke({"poem": poem, "language": "英语"})
    print(f"【{llm.model_name}】的翻译结果")
    print(f"{translation}")
    print("----------------------------------------------")
